import os.path
import _pickle
import pickle
import time

import cv2
import numpy as np

fileList = ['1-1-Conan Gore Fly', '1-2-Front', '1-3-360 Google Spotlight Stories_ HELP',
            '1-4-Conan Weird Al', '1-5-TahitiSurf', '1-6-Falluja',
            '1-7-Cooking Battle', '1-8-Football', '1-9-Rhinos']
salFileList = ["saliency_ds2_topic0", "saliency_ds2_topic1", "saliency_ds2_topic2",
               "saliency_ds2_topic3", "saliency_ds2_topic4", "saliency_ds2_topic5",
               "saliency_ds2_topic6", "saliency_ds2_topic7", "saliency_ds2_topic8"] # attention数据集
videoName = ['Conan1', 'Skiing', 'Alien', 'Conan2', 'Surfing', 'War', 'Cooking', 'Football', 'Rhinos']

frame_path = r"D:\VR_project\saliency-convnet-7.31\framed"
time_path = r"D:\VR_project\saliency-convnet-7.31\saliency_time"
sal_path = r"D:\VR_project\PanoSaliency\data"

if not os.path.exists("timestamp"):
    os.mkdir("timestamp")
if not os.path.exists("saliency"):
    os.mkdir("saliency")

for i,name in enumerate(videoName):
    if i != 8:
        continue
    time_array = np.loadtxt(os.path.join(time_path, f'[{name}]saliency_time.txt'))
    frame_list = os.listdir(os.path.join(frame_path, fileList[i]))
    try:
        all_array = np.array(
            pickle.load(open(os.path.join(sal_path, salFileList[i]), 'rb'), encoding='bytes'), dtype=object)
    except _pickle.UnpicklingError:
        all_array = np.load(os.path.join(sal_path, salFileList[i]), allow_pickle=True)
    saliency_array = np.array([item[2] for item in all_array])

    if not os.path.exists(fileList[i]):
        os.mkdir(fileList[i])
    else:
        for item in os.listdir(fileList[i]):
            os.remove(fileList[i] + '/' + item)

    frameId = 1
    step = 4
    dsize = (160, 90)
    length = len(time_array)
    time_list = []
    sal_list = []

    start = time.time()
    for idx in range(0, length - 20, step):
        frame = cv2.imread(os.path.join(frame_path, f'{fileList[i]}\\{idx+1}.jpg'))
        cv2.imwrite(f'{fileList[i]}/{frameId}.jpg', cv2.resize(frame, dsize))
        time_list.append(round(time_array[idx], 2))
        sal_list.append(saliency_array[idx])
        frameId += 1

    assert frameId-1 == len(time_list) == len(sal_list), "Frame nums is not equal to time_stamp!"
    np.savetxt(f"timestamp/{fileList[i]}.txt", np.array(time_list))
    np.save(f"saliency/{fileList[i]}.npy", np.array(sal_list))
    end = time.time()
    print(f"Sample {frameId-1} frames for {name} finished! Time={round(end-start, 4)}s")
